import gradio as gr
import numpy as np
from PIL import Image
from pinecone import Pinecone
import time
from sklearn.datasets import load_digits

print("=== Pinecone Web应用 ===")

# 初始化Pinecone
def setup_pinecone():
    api_key = 'pcsk_3Hx5py_TFFCuUMU2bVyo3V9yQy4tzpMuotrB4irA2brsi3nncJRkN1Ut7CHUjFm1x8oXn9'
    pc = Pinecone(api_key=api_key)
    return pc.Index("mnist-handwritten-digits")

try:
    index = setup_pinecone()
    print("Pinecone索引连接成功")
except Exception as e:
    print(f"Pinecone连接失败: {e}")
    print("请确保已运行 pinecone_train.py 创建索引")

def preprocess_image(image):
    """
    预处理手写数字图像，使其与MNIST数据集格式兼容
    """
    try:
        # 将图像转换为PIL Image对象
        if isinstance(image, str):
            image = Image.open(image)
        elif isinstance(image, np.ndarray):
            image = Image.fromarray((image * 255).astype('uint8'))
        
        # 将图像转换为灰度
        if image.mode != 'L':
            image = image.convert('L')
        
        # 调整大小为8x8像素（与MNIST数据集相同）
        image = image.resize((8, 8), Image.LANCZOS)
        
        # 转换为numpy数组
        img_array = np.array(image)
        
        # 反转颜色（MNIST是白字黑底，而手写板通常是黑字白底）
        img_array = 255 - img_array
        
        # 归一化到0-16范围（与原始MNIST数据集一致）
        img_array = img_array / 255.0 * 16
        img_array = np.clip(img_array, 0, 16)
        
        # 展平为1D数组
        img_flat = img_array.flatten()
        
        return img_flat.astype(float).tolist()
    
    except Exception as e:
        print(f"图像预处理错误: {e}")
        return None

def predict_digit_pinecone(image, k_value=11):
    """
    使用Pinecone云服务预测手写数字
    """
    if image is None:
        return "请绘制一个数字", ""
    
    try:
        start_time = time.time()
        
        # 预处理图像
        query_vector = preprocess_image(image)
        
        if query_vector is None:
            return "图像处理失败，请重试", ""
        
        # 使用Pinecone查询
        results = index.query(
            vector=query_vector,
            top_k=int(k_value),
            include_metadata=True
        )
        
        query_time = time.time() - start_time
        
        # 统计最近邻的标签
        neighbor_labels = [match['metadata']['label'] for match in results['matches']]
        
        # 使用多数投票决定预测标签
        predicted_label = max(set(neighbor_labels), key=neighbor_labels.count)
        
        # 计算置信度（基于投票比例）
        confidence = neighbor_labels.count(predicted_label) / len(neighbor_labels)
        
        # 构建详细结果
        result = f"预测数字: {predicted_label}\n置信度: {confidence:.2%}\n查询时间: {query_time:.2f}秒\n使用k值: {k_value}"
        
        # 显示最近邻的详细信息
        details = "最近邻结果:\n"
        for i, match in enumerate(results['matches'][:5]):  # 显示前5个结果
            details += f"{i+1}. 数字 {match['metadata']['label']} (相似度: {match['score']:.3f})\n"
        
        return result, details
    
    except Exception as e:
        return f"预测错误: {str(e)}", ""

def create_pinecone_webapp():
    """创建Pinecone Web应用界面"""
    
    with gr.Blocks(title="手写数字识别 - Pinecone KNN") as demo:
        gr.Markdown("# 手写数字识别 - Pinecone KNN")
        gr.Markdown("使用Pinecone云服务进行KNN手写数字识别")
        
        with gr.Row():
            with gr.Column():
                sketchpad = gr.Sketchpad(
                    label="在手写板上绘制数字（0-9）",
                    shape=(200, 200),
                    brush_radius=5
                )
                
                k_slider = gr.Slider(
                    minimum=1,
                    maximum=20,
                    value=11,
                    step=1,
                    label="K值 (最近邻数量)"
                )
                
                predict_btn = gr.Button("识别数字", variant="primary")
                
            with gr.Column():
                result_output = gr.Textbox(
                    label="预测结果",
                    lines=5
                )
                
                details_output = gr.Textbox(
                    label="详细结果",
                    lines=6
                )
        
        # 示例
        gr.Markdown("### 使用说明")
        gr.Markdown("""
        1. 在左侧画板上绘制数字0-9
        2. 调整K值滑块选择最近邻数量
        3. 点击'识别数字'按钮或等待自动识别
        4. 查看右侧的预测结果和详细分析
        
        **注意**: 第一次推理可能会较慢，因为需要初始化云连接
        """)
        
        # 绑定事件
        predict_btn.click(
            fn=predict_digit_pinecone,
            inputs=[sketchpad, k_slider],
            outputs=[result_output, details_output]
        )
        
        # 实时预测
        sketchpad.change(
            fn=predict_digit_pinecone,
            inputs=[sketchpad, k_slider],
            outputs=[result_output, details_output]
        )
    
    return demo

# 备用方案：简化界面
def create_simple_interface():
    """创建简化界面"""
    return gr.Interface(
        fn=lambda image: predict_digit_pinecone(image, 11)[0],
        inputs=gr.Sketchpad(shape=(200, 200)),
        outputs="text",
        title="手写数字识别 - Pinecone",
        description="使用Pinecone云服务识别手写数字",
        live=True
    )

if __name__ == "__main__":
    print("启动Pinecone Web应用...")
    
    try:
        # 验证Pinecone连接
        stats = index.describe_index_stats()
        print("Pinecone索引状态正常")
        print(f"索引统计: {stats}")
        
        # 创建完整应用
        app = create_pinecone_webapp()
        print("请在浏览器中打开显示的本地地址（通常是 http://127.0.0.1:7860）")
        app.launch(share=False, server_port=7860)
        
    except Exception as e:
        print(f"完整应用创建失败: {e}")
        print("尝试简化界面...")
        
        try:
            simple_app = create_simple_interface()
            simple_app.launch(share=False)
        except Exception as e2:
            print(f"简化界面也失败: {e2}")
            print("请检查Pinecone配置和网络连接")